# -*- coding: utf-8 -*-
import random
import numpy as np
import itertools
import ConfigSpace as CS
import ConfigSpace.hyperparameters as CSH
from onnx import TensorProto
from onnx import helper
from . import MappingTypes

class BatchNormalization():
    cs_attributes = CS.ConfigurationSpace()
    attribute_list = ['epsilon', 'momentum', 'training_mode']
    all_attributes = []

    def __init__(self):
        print('BatchNormalization Init--------------------------')
        self.init_all_attributes()
        self.init_cs()

    def init_all_attributes(self):
        for i in range(1, len(BatchNormalization.attribute_list)+1) :
            #此处取所有组合的可能值，存在列表里
            for attribute in itertools.combinations(BatchNormalization.attribute_list, i):
                l = list(attribute)
                BatchNormalization.all_attributes.append(l)

    def init_cs(self):
        length = len(BatchNormalization.all_attributes)
        attributes = CSH.UniformIntegerHyperparameter(name='combination', lower=0, upper=2*length)
        BatchNormalization.cs_attributes.add_hyperparameter(attributes)            

        print('all attributes:')
        for a in BatchNormalization.all_attributes:
            print(a)

    def get_random_combination(self):
        s = BatchNormalization.cs_attributes.sample_configuration()
        index = s.get_dictionary()['combination']
        print('get combination: ', index)

        length = len(BatchNormalization.all_attributes)

        #保证不带任何属性的case可以被覆盖到
        if index >= length:
            print('test no attributes')
            return []

        return BatchNormalization.all_attributes[index]  

    def get_max_attributes_combination(self):
        return len(BatchNormalization.all_attributes)                     

    def get_attributes_by_index_or_random(self, shapes, index):

        print('get_attributes_by_index_or_random, index: ', index)    
            
        attributes = {}
        shape_out = []

        print('arg, shapes:', shapes[0])

        if index != -1:
            attr_combination = BatchNormalization.all_attributes[index]
        else:
            attr_combination = self.get_random_combination()

        if len(attr_combination) == 0:
            return attributes, shape_out    

        if 'epsilon' in attr_combination:
            epsilon = random.random() * 0.00001
            attributes['epsilon'] = epsilon

        if 'momentum' in attr_combination:
            momentum = random.random()
            attributes['momentum'] = momentum

        if 'training_mode' in attr_combination:
            training_mode = random.randint(0, 1)
            attributes['training_mode'] = training_mode

        '''
        if 'training_mode' not in attr_combination:
            training_mode = 1#random.randint(0, 1)
            attributes['training_mode'] = training_mode   
        '''

        return attributes, shape_out 

    def get_initializer_name_list(self):
            initializer_name_list = ['scale', 'B', 'input_mean', 'input_var']
            return initializer_name_list

    def generate_random_initializer(self, in_shapes, out_shape, dtype):
        np_type = MappingTypes.INPUT_TYPE_TO_NP_TYPE[dtype]

        print('generate_random_initializer, dtype:', dtype)

        scale_shape = [in_shapes[0][1]]
        scale_array = np.random.randn(in_shapes[0][1]).astype(np_type)

        bias_shape = [in_shapes[0][1]]
        bias_array = np.random.randn(in_shapes[0][1]).astype(np_type)

        input_mean_shape = [in_shapes[0][1]]
        input_mean_array = np.random.randn(in_shapes[0][1]).astype(np_type)

        input_var_shape = [in_shapes[0][1]]
        input_var_array = np.random.randn(in_shapes[0][1]).astype(np_type)

        print('scale_shape: ', scale_shape)
        print('scale_array: ', scale_array)

        print('bias_shape: ', bias_shape)
        print('bias_array: ', bias_array)

        print('input_mean_shape: ', input_mean_shape)
        print('input_mean_array: ', input_mean_array)

        print('input_var_shape: ', input_var_shape)
        print('input_var_array: ', input_var_array)

        tensor_type = MappingTypes.get_tensor_type(dtype)

        s = helper.make_tensor('s', tensor_type, scale_shape, scale_array)    
        bias = helper.make_tensor('bias', tensor_type, bias_shape, bias_array)    
        mean = helper.make_tensor('input_mean', tensor_type, input_mean_shape, input_mean_array)  
        var = helper.make_tensor('input_var', tensor_type, input_var_shape, input_var_array)   

        initializer = [s, bias, mean, var]

        return initializer

    def get_input_name(self):
        input_name = ['X', 's', 'bias', 'input_mean', 'input_var']
        return input_name

    def input_shape_correct(self, shapes):
        if len(shapes[0]) < 3:
            for i in range(3 - len(shapes[0])):
                shapes[0].append(random.randint(2, 8))

        return True

    def get_output_tensor(self, in_shapes, dtype, attributes):
        output_tensor_list = []

        print('get_output_tensor, in_shapes:', in_shapes, ', dtype:', dtype, ', attributes:', attributes)

        t = MappingTypes.get_tensor_type(dtype)
        y = helper.make_tensor_value_info('Y', t, in_shapes)

        output_tensor_list.append(y)

        if 'training_mode' in attributes.keys() and attributes['training_mode'] == 1:
            mean = helper.make_tensor_value_info('running_mean', t, [in_shapes[1]])
            output_tensor_list.append(mean)

            var = helper.make_tensor_value_info('running_var', t, [in_shapes[1]])
            output_tensor_list.append(var)

        return output_tensor_list

    def get_output_name(self, attributes):
        output_name_list = []

        print('get_output_tensor,  attributes:', attributes)

        output_name_list.append('Y')

        if 'training_mode' in attributes.keys() and attributes['training_mode'] == 1:
            output_name_list.append('running_mean')
            output_name_list.append('running_var')

        return output_name_list        








